import requests
import json
from common.ploto_enum import TrainingJobEnum
from conf.conf import obs_endpoint, ces_project_id, IAMDomain, IAMUser, IAMPassword, IAMproject

region = obs_endpoint.split('.')[1]

dataset_api_domain = "https://modelarts.{}.myhuaweicloud.com/v2/{}/datasets/".format(region, ces_project_id)
algorithms_api_domain = "https://modelarts.{}.myhuaweicloud.com/v2/{}/algorithms/".format(region, ces_project_id)
training_job_api_domain = "https://modelarts.{}.myhuaweicloud.com/v2/{}/training-jobs/".format(region, ces_project_id)
training_job_search_api_domain = "https://modelarts.{}.myhuaweicloud.com/v2/{}/training-job-searches/".\
    format(region, ces_project_id)
models_api_domain = "https://modelarts.{}.myhuaweicloud.com/v1/{}/models/".format(region, ces_project_id)


class DataSet:
    def __init__(self) -> None:
        self.info = {}
        self.token = iam_auth()
        self.header = {"Content-Type": "application/json", "X-Auth-Token": self.token}
        self.info['headers'] = self.header

    def get_datasets(self, page, page_size, name):
        # 获取数据集列表
        url = dataset_api_domain[:-1]
        offset = (page - 1) * page_size
        request_url = "{}?offset={}&limit={}".format(url, offset, page_size)
        if name != "":
            request_url = "{}&search_content={}".format(request_url, name)
        resp = requests.get(request_url, **self.info)
        return resp

    def get_dataset(self, dataset_id):
        # 获取数据集详细信息
        request_url = dataset_api_domain + dataset_id
        resp = requests.get(request_url, **self.info)
        return resp

    def add_dataset(self, param):
        # 创建数据集
        data_path, name, description, work_path, dataset_type = param
        request_url = dataset_api_domain[:-1]
        body = {
            "data_sources": [{
                "data_path": data_path
            }],
            "import_data": True,
            "dataset_name": name,
            "dataset_type": dataset_type,
            "description": description,
            "work_path": work_path,
            "labels": [{
                "name": "Car",
                "type": 1
            }]
        }
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def sysn_data(self, dataset_id):
        # 同步数据
        request_url = dataset_api_domain + dataset_id + '/sync-data'
        resp = requests.post(request_url, **self.info)
        return resp

    def sysn_data_status(self, dataset_id):
        # 查询同步数据状态
        request_url = dataset_api_domain + dataset_id + '/sync-data/status'
        resp = requests.get(request_url, **self.info)
        return resp

    def get_import_data_info(self, dataset_id):
        # 查询导入数据任务状态
        request_url = dataset_api_domain + dataset_id + '/import-tasks'
        resp = requests.get(request_url, **self.info)
        return resp

    def import_task(self, dataset_id, data_url):
        # 导入数据
        request_url = dataset_api_domain + dataset_id + '/import-tasks'
        body = {"import_path": data_url}
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def auto_task(self, dataset_id, algorithm_type):
        # 启动智能标注--主动学习
        request_url = dataset_api_domain + dataset_id + '/tasks'
        body = {
            "task_type": "auto-label",
            "collect_key_sample": True,
            "config": {
                "algorithm_type": algorithm_type
            }
        }
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def pre_task(self, dataset_id, model_id, specification, instance_count):
        # 启动智能标注--预标注
        request_url = dataset_api_domain + dataset_id + '/tasks'
        body = {
            "task_type": "pre-label",
            "model_id": model_id,
            "collect_key_sample": True,
            "config": {
                "inf_config_list": [{
                    "specification": specification,
                    "instance_count": instance_count
                }]
            }
        }
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def update_dataset(self, dataset_id, dataset_name, description):
        # 修改数据集
        request_url = dataset_api_domain + dataset_id
        body = {
            "dataset_name": dataset_name,
            "description": description
        }
        resp = requests.put(request_url, **self.info, data=json.dumps(body))
        return resp

    def delete_dataset(self, dataset_id):
        # 删除数据集
        request_url = dataset_api_domain + dataset_id
        resp = requests.delete(request_url, **self.info)
        return resp


class Algorithm:
    def __init__(self) -> None:
        self.info = {}
        self.token = iam_auth()
        self.header = {"Content-Type": "application/json", "X-Auth-Token": self.token}
        self.info['headers'] = self.header

    def get_algorithms(self, page, page_size, name):
        # 查询算法列表
        url = algorithms_api_domain[:-1]
        offset = (page - 1) * page_size
        request_url = "{}?offset={}&limit={}".format(url, offset, page_size)
        if name != "":
            request_url = "{}&searches=name%3A{}".format(request_url, name)
        resp = requests.get(request_url, **self.info)
        return resp

    def get_hub_algorithms(self):
        # 查询订阅算法(AI hub)
        request_url = "https://modelarts.{}.myhuaweicloud.com/v1/aihub/subscriptions?content_types=algo".format(region)
        resp = requests.get(request_url, **self.info)
        return resp

    def create_algorithms(self, algorithms_detail):
        # 创建算法
        request_url = algorithms_api_domain[:-1]
        body = {}
        algo_type = int(algorithms_detail.type)
        # 预置框架
        if algo_type == TrainingJobEnum.FRAME.code:
            body = {
                "metadata": {
                    "name": algorithms_detail.name,
                    "description": algorithms_detail.description
                },
                "job_config": {
                    "code_dir": algorithms_detail.code_dir,
                    "boot_file": algorithms_detail.boot_file,
                    "engine": {
                        "engine_name": algorithms_detail.engine_name,
                        "engine_version": algorithms_detail.engine_version
                    }
                }
            }
        # 自定义
        else:
            body = {
                "metadata": {
                    "name": algorithms_detail.name,
                    "description": algorithms_detail.description
                },
                "job_config": {
                    "command": algorithms_detail.command,
                    "engine": {
                        "image_url": algorithms_detail.image_url
                    }
                }
            }
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def update_algorithms(self, algorithm_id, algorithm_name, description):
        # 更新算法
        request_url = algorithms_api_domain + algorithm_id
        details = requests.get(request_url, **self.info)
        text = json.loads(details.text)
        engine_id = text["job_config"]["engine"]["engine_id"]
        engine_name = text["job_config"]["engine"]["engine_name"]
        engine_version = text["job_config"]["engine"]["engine_version"]
        code_dir = text["job_config"]["code_dir"]
        boot_file = text["job_config"]["boot_file"]
        body = {
            "metadata": {
                "name": algorithm_name,
                "description": description
            },
            "job_config": {
                "engine": {
                    "engine_id": engine_id,
                    "engine_name": engine_name,
                    "engine_version": engine_version
                },
                "code_dir": code_dir,
                "boot_file": boot_file
            }
        }
        resp = requests.put(request_url, **self.info, data=json.dumps(body))
        return resp

    def delete_algorithms(self, algorithm_id):
        # 删除算法
        request_url = algorithms_api_domain + algorithm_id
        resp = requests.delete(request_url, **self.info)
        return resp


class TrainingJob:
    def __init__(self) -> None:
        self.info = {}
        self.token = iam_auth()
        self.header = {"Content-Type": "application/json", "X-Auth-Token": self.token}
        self.info['headers'] = self.header

    def get_training_job(self, page, page_size, name):
        # 查询任务列表 api请求方式为post
        url = training_job_search_api_domain[:-1]
        offset = (page - 1) * page_size
        if name != "":
            body = {
                "offset": offset,
                "limit": page_size,
                "filters": [{
                    "key": "name",
                    "operator": "like",
                    "value": [name]
                }]
            }
            resp = requests.post(url, **self.info, data=json.dumps(body))
        else:
            request_url = "{}?offset={}&limit={}".format(url, offset, page_size)
            resp = requests.post(request_url, **self.info)
        return resp

    def create_training_jobs(self, job_detail, create_type):
        # 创建训练任务
        request_url = training_job_api_domain[:-1]
        body = {}
        create_type = int(create_type)
        if hasattr(job_detail, "node_count"):
            node_count = int(job_detail.node_count)
        else:
            node_count = 1
        if create_type == TrainingJobEnum.FRAME.code:
            # 自定义算法--预置框架
            body = {
                "kind": "job",
                "metadata": {
                    "name": job_detail.name,
                    "description": job_detail.description
                },
                "algorithm": {
                    "code_dir": job_detail.code_dir,
                    "boot_file": job_detail.boot_file,
                    "engine": {
                        "engine_name": job_detail.engine_name,
                        "engine_version": job_detail.engine_version
                    }
                },
                "spec": {
                    "resource": {
                        "flavor_id": job_detail.flavor_id,
                        "node_count": node_count
                    }
                }
            }
        if create_type == TrainingJobEnum.CUSTOM.code:
            # 自定义算法--自定义
            body = {
                "kind": "job",
                "metadata": {
                    "name": job_detail.name,
                    "description": job_detail.description
                },
                "algorithm": {
                    "command": job_detail.command,
                    "engine": {
                        "image_url": job_detail.image_url
                    }
                },
                "spec": {
                    "resource": {
                        "flavor_id": job_detail.flavor_id,
                        "node_count": node_count
                    }
                }
            }
        if create_type == TrainingJobEnum.ALGORITHM.code:
            # 我的算法
            body = {
                "kind": "job",
                "metadata": {
                    "name": job_detail.name,
                    "description": job_detail.description
                },
                "algorithm": {
                    "id": job_detail.id
                },
                "spec": {
                    "resource": {
                        "flavor_id": job_detail.flavor_id,
                        "node_count": node_count
                    }
                }
            }
        if create_type == TrainingJobEnum.AIHUB.code:
            body = {
                "kind": "job",
                "metadata": {
                    "name": job_detail.name,
                    "description": job_detail.description
                },
                "algorithm": {
                    "subscription_id": job_detail.subscription_id,
                    "item_version_id": job_detail.item_version_id
                },
                "spec": {
                    "resource": {
                        "flavor_id": job_detail.flavor_id,
                        "node_count": node_count
                    }
                }
            }
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def delete_training_job(self, job_id):
        # 删除训练任务
        request_url = training_job_api_domain + job_id
        resp = requests.delete(request_url, **self.info)
        return resp

    def stop_training_job(self, job_id):
        # 终止训练任务
        request_url = training_job_api_domain + job_id + '/actions'
        body = {
            "action_type": "terminate"
        }
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def update_training_job(self, job_id, description):
        # 修改训练任务描述
        request_url = training_job_api_domain + job_id
        body = {
            "description": description
        }
        resp = requests.put(request_url, **self.info, data=json.dumps(body))
        return resp


class Model:
    def __init__(self) -> None:
        self.info = {}
        self.token = iam_auth()
        self.header = {"Content-Type": "application/json", "X-Auth-Token": self.token}
        self.info['headers'] = self.header

    def get_models(self, page, page_size, name):
        # 查询模型列表
        url = models_api_domain[:-1]
        offset = (page - 1) * page_size
        request_url = "{}?offset={}&limit={}".format(url, offset, page_size)
        if name != "":
            if name.find('_'):
                request_url = "{}&model_name={}&exact_match=1".format(request_url, name)
            else:
                request_url = "{}&model_name={}".format(request_url, name)
        resp = requests.get(request_url, **self.info)
        return resp

    def get_hub_models(self):
        # 查询订阅算法(AI hub)
        request_url = "https://modelarts.{}.myhuaweicloud.com/v1/aihub/subscriptions?content_types=model".format(region)
        resp = requests.get(request_url, **self.info)
        return resp

    def create_model(self, detail):
        # 导入模型
        request_url = models_api_domain[:-1]
        body = {
            "model_name": detail.model_name,
            "model_version": detail.model_version,
            "source_location": detail.source_location,
            "model_type": detail.model_type,
            "runtime": detail.runtime
        }
        resp = requests.post(request_url, **self.info, data=json.dumps(body))
        return resp

    def delete_model(self, model_id):
        # 删除模型
        request_url = models_api_domain + model_id
        resp = requests.delete(request_url, **self.info)
        return resp


class Resource:
    def __init__(self) -> None:
        self.info = {}
        self.token = iam_auth()
        self.header = {"Content-Type": "application/json", "X-Auth-Token": self.token}
        self.info['headers'] = self.header

    def get_engines(self):
        # 获取AI引擎列表
        request_url = "https://modelarts.{}.myhuaweicloud.com/v2/{}/training-job-engines".format(region, ces_project_id)
        resp = requests.get(request_url, **self.info)
        return resp

    def get_flavors(self):
        # 获取训练资源规格列表
        request_url = "https://modelarts.{}.myhuaweicloud.com/v2/{}/training-job-flavors".format(region, ces_project_id)
        resp = requests.get(request_url, **self.info)
        return resp


IAM_info = {
    "auth": {
        "identity": {
            "methods": [
                "password"
            ],
            "password": {
                "user": {
                    "domain": {
                        "name": "%s"
                    },
                    "name": "%s",
                    "password": "%s"
                }
            }
        },
        "scope": {
            "project": {
                "name": "%s"
            }
        }
    }
}


def iam_auth():
    iam_info_json = json.dumps(IAM_info)
    iam_info_json = iam_info_json % (IAMDomain, IAMUser, IAMPassword, IAMproject)
    auth_api = "https://iam.{}.myhuaweicloud.com/v3/auth/tokens".format(IAMproject)
    headers = {"Content-Type": "application/json"}
    auth_result = requests.post(auth_api, data=iam_info_json, headers=headers)
    token = auth_result.headers['X-Subject-Token']
    headers['X-Auth-Token'] = token
    return token
